# -*- coding: UTF-8 -*-

# %% [markdown]
#
import numpy as np
import matplotlib.pyplot as plt
import control as ct
import sympy as sp

def ztrans2diffexpr(tf_z: sp.core.mul.Mul):
    degree_num = sp.degree(sp.numer(tf_z), z)
    degree_den = sp.degree(sp.denom(tf_z), z)
    tf_z_num = sp.numer(tf_z)  # /z**degree_num
    tf_z_num = sp.simplify(tf_z_num)
    tf_z_num_coeff = [tf_z_num.coeff(z, i)
                                for i in range(degree_num+1)]
    # tf_z_num_coeff.reverse()
    tf_z_den = sp.denom(tf_z)  # /z**degree_den
    tf_z_den = sp.simplify(tf_z_den)
    tf_z_den_coeff = [tf_z_den.coeff(z, i)
                                for i in range(degree_den+1)]
    # tf_z_den_coeff.reverse()

    k = tf_z_den_coeff[-1]
    # m_ = [x/k for x in tf_z_num_coeff]
    # n_ = [-x/k for x in tf_z_den_coeff[:-1]]
    # m_.reverse()
    # n_.reverse()
    m = (np.array(tf_z_num_coeff)/k)[::-1]
    n = (-np.array(tf_z_den_coeff)[:-1]/k)[::-1]

    return m, n


def notch_filter_generate_gain(Ts_, d_deepth_, f_width_, f0_):

    d_deepth, f_width, f0 = sp.symbols('d_deepth f_width f0')
    Ts = sp.symbols('Ts')
    # d_deepth, f_width, f0 = d_deepth_, f_width_, f0_
    # Ts = Ts_

    w0 = 2*sp.pi*f0
    w_width = 2*sp.pi*f_width*1.5
    # w_width = (sp.sqrt((4*sp.pi*f_width/w0)**2+4)+1)/(-4*d_deepth**2+2)*(2*w0)
    
    s = sp.symbols('s')
    z = sp.symbols('z')

    # s = (1-z**-1)/Ts # backward
    # s = 2/Ts*(z-1)/(z+1) # backward

    sys = (s**2 + w_width*d_deepth*s + w0**2)/(s**2 + w_width*s + w0**2)
    sys = sys.subs(s, 2/Ts*(z-1)/(z+1))
    sys_tmp = sys

    sys_tmp = sp.simplify(sys_tmp)
    sys_tmp = sp.expand(sys_tmp, frac=True)
    sys_tmp = sp.collect(sys_tmp, z)
    sys_tmp_value = sys_tmp
    sys_tmp_value = sys_tmp.subs({Ts:Ts_, d_deepth:d_deepth_, f_width:f_width_, f0:f0_})
    # sys_tmp_value = sys_tmp.subs({Ts:1e-4, d_deepth:0.5, f_width:3, f0:15})

    m, n = ztrans2diffexpr(sys_tmp_value)
    return m, n
    
def notch_filter(x, x_last1, x_last2, y_last1, y_last2, m, n):

    y = (m[0]*x + m[1]*x_last1 + m[2]*x_last2) + \
            (n[0]*y_last1 + n[1]*y_last2)

    return y

if __name__ == "__main__":
    pass
    # %% [markdown]
    #
    x, y = sp.symbols('x y')
    z = sp.symbols('z')

    # %% [markdown]
    # form
    Ts, d_deepth, f_width, f0 = sp.symbols('Ts d_deepth f_width f0')
    m, n = notch_filter_generate_gain(Ts, d_deepth, f_width, f0)
    print('m[0]:{}'.format(m[0]))
    print('m[1]:{}'.format(m[1]))
    print('m[2]:{}'.format(m[2]))
    print('n[0]:{}'.format(n[0]))
    print('n[1]:{}'.format(n[1]))
    
    
    # %% [markdown]
    # data test
    Ts0, d_deepth0, f_width0, f00 = sp.symbols('Ts0 d_deepth0 f_width0 f00')
    # m, n = notch_filter_generate_gain(Ts0, d_deepth0, f_width0, f00)
    m, n = notch_filter_generate_gain(1e-4, 0.5, 30, 60)
    m = [float(m[i]) for i in range(len(m))]
    n = [float(n[i]) for i in range(len(n))]
    print(m)
    print(n)

    x, x_last1, x_last2, y_last1, y_last2 = 0, 0, 0, 0, 0

    f = 45
    t = np.arange(20000)*1e-4
    sin_data = np.sin(2*np.pi*f*t)
    # plt.scatter(t, sin_data)
    # plt.show()

    x_list = []
    y_list = []
    y = 0
    for i in range(len(t)):
        x_last2 = x_last1
        x_last1 = x
        x = sin_data[i]
        y_last2 = y_last1
        y_last1 = y
        x_list.append(x)
        y = notch_filter(x, x_last1, x_last2, y_last1, y_last2, m, n)
        y_list.append(y)

    plt.scatter(t, y_list)
    plt.show()
    print(max(x_list[10000:]))
    print(max(y_list[10000:]))